package pocx

import (
	"context"
	"encoding/json"
	"fmt"
	"gitee.com/wudicaidou/menciis-pocx/pocbase"
	"github.com/antlabs/strsim"
	"strconv"
	"strings"
	"time"
)

type RpcPoc struct {
	Id     string   `json:"id" #:"id 信息"`
	Name   string   `json:"name" #:"插件名称"`
	Tags   []string `json:"tags"`
	Data   []byte   `json:"data"`
	config *Config
}

func (t *RpcPoc) Init(ctx context.Context) error {
	config, err := ContextConfig(ctx)
	if err != nil {
		return err
	}
	t.config = config
	if t.config.RpcClient == nil {
		return fmt.Errorf("rpcClient need to be configured")
	}
	if t.Id == "" {
		t.getInfo()
	}
	return nil
}

func (t *RpcPoc) getInfo() {
	serialize, err := t.config.RpcClient.CheckPythonSerialize(t.Data)
	err = json.Unmarshal(*serialize, t)
	if t.Name == "" {
		t.Name = t.Id
	}
	if err != nil {
		return
	}
}

func (t *RpcPoc) Meta() *pocbase.MetaInfo {
	return &pocbase.MetaInfo{
		Plugin: &pocbase.PluginInfo{
			ID:   t.Id,
			Name: t.Name,
		},
	}
}
func (t *RpcPoc) Match(ctx context.Context, in []pocbase.Fingerprint) bool {
	if len(in) == 0 || len(t.Tags) == 0 {
		return true
	} else {
		if len(in) > 0 {
			for _, v := range in {
				for _, y := range t.Tags {
					if strsim.Compare(strings.ToLower(v.Product), strings.ToLower(y)) > 0.3 {
						return true
					}
				}
			}
		}
		return false
	}
}

func (t *RpcPoc) Scan(ctx context.Context, in interface{}) (pocbase.SnapShot, pocbase.Event, error) {
	var snapShot pocbase.SnapShot
	if t.config.RpcClient == nil {
		return snapShot, nil, fmt.Errorf("rpcClient need to be configured")
	}
	if t.Id == "" {
		return snapShot, nil, fmt.Errorf("poc need id")
	}
	srv, isService := in.(*pocbase.ServiceAsset)
	requestAsset, isHTTP := in.(*pocbase.RequestAsset)
	var host, port string
	if isHTTP {
		req := requestAsset.Req
		host = req.GetHostName()
		port = req.GetPort()
		if port == "" {
			if req.GetScheme() == "https" {
				port = "443"
			} else {
				port = "80"
			}
		}
	} else if isService {
		host = srv.Host
		port = strconv.Itoa(srv.Port)
	} else {
		return snapShot, nil, fmt.Errorf("pocbase.RequestAsset is of the wrong type")
	}
	if t.Id == "" {
		t.Id = t.Name
	}

	result, err := t.config.RpcClient.ExecPythonSerialize(host, port, t.Id, t.Data)
	if err != nil {
		return snapShot, nil, err
	}
	event, err := t.getResult(result, host, port)
	return snapShot, event, err
}

type Result struct {
	VulID    string `json:"vulID"`
	Name     string `json:"name"`
	Desc     string `json:"desc"`
	Url      string `json:"url"`
	Response string `json:"response"`
	Request  string `json:"request"`
	Info     string `json:"info"`
}

type pythonResult struct {
	Result Result `json:"result"`
	PyErr  string `json:"err"`
}

// 获取返回内容
func (t *RpcPoc) getResult(data *[]byte, host string, port string) (pocbase.Event, error) {
	var target *pocbase.TargetInfo
	var detail *pocbase.VulDetail
	var result pythonResult
	// todo 执行失败要做处理
	err := json.Unmarshal(*data, &result)
	if err != nil {
		return nil, err
	} else if result.PyErr != "" {
		return nil, fmt.Errorf(result.PyErr)
	}
	info := result.Result
	if info.VulID == "" && info.Name == "" && info.Desc == "" && info.Url == "" && info.Info == "" {
		return nil, nil
	}
	if t.Name == "" {
		t.Name = info.VulID
	}
	if t.Name == "" {
		t.Name = info.Name
	}
	iPort, err := strconv.Atoi(port)
	if err != nil {
		return nil, err
	}
	target = &pocbase.TargetInfo{
		Host: host,
		Path: info.Url,
		Port: iPort,
	}
	var content [][2][]byte
	content = append(content, [2][]byte{[]byte(info.Request), []byte(info.Response)})
	detail = &pocbase.VulDetail{
		Payload:       "",
		SnapShot:      nil,
		ExtractedInfo: map[string]interface{}{},
	}
	pluginInfo := &pocbase.PluginInfo{
		ID:   t.Id,
		Name: t.Name,
	}
	resultEvent := &pocbase.VulEvent{
		Timestamp: time.Now(),
		Plugin:    pluginInfo,
		Details:   detail,
		Target:    target,
	}

	return resultEvent, nil
}

func (t *RpcPoc) Close() {}

func (t *RpcPoc) SetProxy(proxy string) {}
